"""
安装FlagEmbedding: pip install FlagEmbedding
"""
from FlagEmbedding import FlagModel
sentences_1 = ["样例数据-1", "样例数据-2"] # 两个句子信息
sentences_2 = ["样例数据-3", "样例数据-4"] # 两个句子信息
model = FlagModel('./data/BAAI/bge-small-zh',
                  use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
embeddings_1 = model.encode(sentences_1)
embeddings_2 = model.encode(sentences_2)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

# for s2p(short query to long passage) retrieval task, suggest to use encode_queries() which will automatically add the instruction to each query
# corpus in retrieval task can still use encode() or encode_corpus(), since they don't need instruction
# queries = ['query_1', 'query_2']
# passages = ["样例文档-1", "样例文档-2"]
# q_embeddings = model.encode_queries(queries)
# p_embeddings = model.encode(passages)
# scores = q_embeddings @ p_embeddings.T